Coverage for src/quickmcp/factory/wrappers.py: 79%

1""" 

2Tool wrappers for converting Python functions to MCP tools. 

3""" 

4 

5import asyncio 

6import inspect 

7import logging 

8from typing import Any, Callable, Dict, get_type_hints, Optional 

9from abc import ABC, abstractmethod 

10 

11from .config import FactoryConfig, DEFAULT_CONFIG 

12from .type_conversion import TypeConverter 

13from .errors import ToolRegistrationError, TypeConversionError, log_factory_error 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18class ToolWrapper(ABC): 

19 """Abstract base class for tool wrappers.""" 

20 

21 def __init__(self, func: Callable, name: str, config: Optional[FactoryConfig] = None): 

22 self.func = func 

23 self.name = name 

24 self.config = config or DEFAULT_CONFIG 

25 self.type_converter = TypeConverter(config) 

26 

27 # Extract function metadata 

28 self.signature = inspect.signature(func) 

29 self.doc = func.__doc__ or f"Execute {name}" 

30 self.type_hints = self._get_type_hints_safely() 

31 

32 # Create the wrapper function 

33 self.wrapper = self._create_wrapper() 

34 

35 # Copy metadata to wrapper 

36 self.wrapper.__name__ = name 

37 self.wrapper.__doc__ = self.doc 

38 self.wrapper.__annotations__ = getattr(func, '__annotations__', {}) 

39 

40 def _get_type_hints_safely(self) -> Dict[str, Any]: 

41 """Get type hints with error handling.""" 

42 try: 

43 if self.config.cache_type_hints: 

44 # In a real implementation, you'd cache this 

45 pass 

46 return get_type_hints(self.func) 

47 except NameError as e: 

48 logger.debug(f"Type hints unavailable for {self.name}: {e}") 

49 return {} 

50 except ModuleNotFoundError as e: 

51 logger.debug(f"Module not found for type hints in {self.name}: {e}") 

52 return {} 

53 except Exception as e: 

54 logger.warning(f"Unexpected error getting type hints for {self.name}: {e}") 

55 return {} 

56 

57 @abstractmethod 

58 def _create_wrapper(self) -> Callable: 

59 """Create the actual wrapper function.""" 

60 pass 

61 

62 def _convert_arguments(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: 

63 """Convert arguments to appropriate types.""" 

64 converted_args = {} 

65 

66 for param_name, param in self.signature.parameters.items(): 

67 if param_name in kwargs: 

68 value = kwargs[param_name] 

69 

70 # Try to convert based on type hints 

71 if param_name in self.type_hints: 

72 try: 

73 target_type = self.type_hints[param_name] 

74 converted_value = self.type_converter.convert_value(value, target_type, param_name) 

75 converted_args[param_name] = converted_value 

76 except TypeConversionError as e: 

77 if self.config.strict_type_conversion: 

78 raise ToolRegistrationError( 

79 f"Type conversion failed for parameter '{param_name}' in tool '{self.name}': {e}", 

80 self.name, e 

81 ) 

82 else: 

83 logger.warning(f"Type conversion failed for {param_name}, using original value: {e}") 

84 converted_args[param_name] = value 

85 else: 

86 # No type hint, use as-is 

87 converted_args[param_name] = value 

88 

89 elif param.kind == inspect.Parameter.VAR_KEYWORD: 

90 # Handle **kwargs - add any remaining arguments 

91 for k, v in kwargs.items(): 

92 if k not in converted_args: 

93 converted_args[k] = v 

94 

95 return converted_args 

96 

97 def _convert_result(self, result: Any) -> Any: 

98 """Convert result to JSON-serializable format.""" 

99 if result is None: 

100 return {"success": True} 

101 elif isinstance(result, (str, int, float, bool, list, dict)): 

102 return result 

103 elif hasattr(result, '__dict__'): 

104 # Try to convert to dict 

105 try: 

106 return result.__dict__ 

107 except Exception as e: 

108 logger.debug(f"Failed to convert result to dict: {e}") 

109 return str(result) 

110 else: 

111 return str(result) 

112 

113 

114class SyncToolWrapper(ToolWrapper): 

115 """Wrapper for synchronous functions.""" 

116 

117 def _create_wrapper(self) -> Callable: 

118 """Create a synchronous wrapper.""" 

119 

120 def sync_wrapper(**kwargs): 

121 try: 

122 # Convert arguments 

123 converted_args = self._convert_arguments(kwargs) 

124 

125 # Call the original function 

126 result = self.func(**converted_args) 

127 

128 # Convert result 

129 return self._convert_result(result) 

130 

131 except Exception as e: 

132 error_msg = f"Error executing tool '{self.name}': {e}" 

133 log_factory_error(e, f"Tool '{self.name}'", logging.ERROR) 

134 

135 if self.config.verbose_errors: 

136 return { 

137 "error": str(e), 

138 "error_type": type(e).__name__, 

139 "tool": self.name 

140 } 

141 else: 

142 return {"error": "Tool execution failed"} 

143 

144 return sync_wrapper 

145 

146 

147class AsyncToolWrapper(ToolWrapper): 

148 """Wrapper for asynchronous functions.""" 

149 

150 def _create_wrapper(self) -> Callable: 

151 """Create an asynchronous wrapper.""" 

152 

153 async def async_wrapper(**kwargs): 

154 try: 

155 # Convert arguments 

156 converted_args = self._convert_arguments(kwargs) 

157 

158 # Call the original async function 

159 result = await self.func(**converted_args) 

160 

161 # Convert result 

162 return self._convert_result(result) 

163 

164 except Exception as e: 

165 error_msg = f"Error executing async tool '{self.name}': {e}" 

166 log_factory_error(e, f"Async tool '{self.name}'", logging.ERROR) 

167 

168 if self.config.verbose_errors: 

169 return { 

170 "error": str(e), 

171 "error_type": type(e).__name__, 

172 "tool": self.name 

173 } 

174 else: 

175 return {"error": "Tool execution failed"} 

176 

177 return async_wrapper 

178 

179 

180class MethodToolWrapper(ToolWrapper): 

181 """Wrapper for instance methods.""" 

182 

183 def __init__(self, method: Callable, instance: Any, name: str, config: Optional[FactoryConfig] = None): 

184 self.instance = instance 

185 super().__init__(method, name, config) 

186 

187 def _create_wrapper(self) -> Callable: 

188 """Create a wrapper that handles the instance method.""" 

189 

190 # Check if the method is async 

191 if inspect.iscoroutinefunction(self.func): 

192 return self._create_async_method_wrapper() 

193 else: 

194 return self._create_sync_method_wrapper() 

195 

196 def _create_sync_method_wrapper(self) -> Callable: 

197 """Create a synchronous method wrapper.""" 

198 

199 def sync_method_wrapper(**kwargs): 

200 try: 

201 # Convert arguments 

202 converted_args = self._convert_arguments(kwargs) 

203 

204 # Call the method on the instance 

205 result = self.func(**converted_args) 

206 

207 # Convert result 

208 return self._convert_result(result) 

209 

210 except Exception as e: 

211 error_msg = f"Error executing method tool '{self.name}': {e}" 

212 log_factory_error(e, f"Method tool '{self.name}'", logging.ERROR) 

213 

214 if self.config.verbose_errors: 

215 return { 

216 "error": str(e), 

217 "error_type": type(e).__name__, 

218 "tool": self.name, 

219 "method": True 

220 } 

221 else: 

222 return {"error": "Method execution failed"} 

223 

224 return sync_method_wrapper 

225 

226 def _create_async_method_wrapper(self) -> Callable: 

227 """Create an asynchronous method wrapper.""" 

228 

229 async def async_method_wrapper(**kwargs): 

230 try: 

231 # Convert arguments 

232 converted_args = self._convert_arguments(kwargs) 

233 

234 # Call the async method on the instance 

235 result = await self.func(**converted_args) 

236 

237 # Convert result 

238 return self._convert_result(result) 

239 

240 except Exception as e: 

241 error_msg = f"Error executing async method tool '{self.name}': {e}" 

242 log_factory_error(e, f"Async method tool '{self.name}'", logging.ERROR) 

243 

244 if self.config.verbose_errors: 

245 return { 

246 "error": str(e), 

247 "error_type": type(e).__name__, 

248 "tool": self.name, 

249 "method": True 

250 } 

251 else: 

252 return {"error": "Method execution failed"} 

253 

254 return async_method_wrapper 

255 

256 

257class ToolWrapperFactory: 

258 """Factory for creating appropriate tool wrappers.""" 

259 

260 def __init__(self, config: Optional[FactoryConfig] = None): 

261 self.config = config or DEFAULT_CONFIG 

262 

263 def create_wrapper(self, func: Callable, name: str) -> ToolWrapper: 

264 """Create an appropriate wrapper for the given function.""" 

265 if inspect.iscoroutinefunction(func): 

266 return AsyncToolWrapper(func, name, self.config) 

267 else: 

268 return SyncToolWrapper(func, name, self.config) 

269 

270 def create_method_wrapper(self, method: Callable, instance: Any, name: str) -> ToolWrapper: 

271 """Create an appropriate wrapper for the given method.""" 

272 return MethodToolWrapper(method, instance, name, self.config) 

273 

274 

275def create_tool_wrapper(func: Callable, name: str, config: Optional[FactoryConfig] = None) -> Callable: 

276 """ 

277 Convenience function to create a tool wrapper. 

278  

279 Args: 

280 func: The function to wrap 

281 name: The name for the tool 

282 config: Optional configuration 

283  

284 Returns: 

285 The wrapped function ready to use as an MCP tool 

286 """ 

287 factory = ToolWrapperFactory(config) 

288 wrapper = factory.create_wrapper(func, name) 

289 return wrapper.wrapper