Coverage for session_buddy / utils / lazy_imports.py: 28.86%

1#!/usr/bin/env python3 

2"""Lazy loading utilities for optional dependencies. 

3 

4This module provides lazy loading for heavy or optional dependencies to improve 

5startup performance and handle missing dependencies gracefully. 

6""" 

7 

8import importlib 

9from collections.abc import Callable 

10from functools import wraps 

11from typing import Any, Never 

12 

13from .logging import get_session_logger 

14 

15# Lazy-load logger to avoid DI initialization issues during imports 

16_logger: Any = None 

17 

18 

19def _get_logger() -> Any: 

20 """Get logger instance, initializing on first use.""" 

21 global _logger 

22 if _logger is None: 

23 try: 

24 _logger = get_session_logger() 

25 except Exception: 

26 # Fallback to basic logging if DI not initialized 

27 import logging 

28 

29 _logger = logging.getLogger(__name__) 

30 return _logger 

31 

32 

33class LazyImport: 

34 """Lazy import wrapper that loads modules on first access.""" 

35 

36 def __init__( 

37 self, 

38 module_name: str, 

39 fallback_value: Any = None, 

40 import_error_msg: str | None = None, 

41 ) -> None: 

42 self.module_name = module_name 

43 self.fallback_value = fallback_value 

44 self.import_error_msg = import_error_msg 

45 self._module: Any = None 

46 self._import_attempted = False 

47 self._import_failed = False 

48 

49 def __getattr__(self, name: str) -> Any: 

50 if not self._import_attempted: 

51 self._try_import() 

52 

53 if self._import_failed: 

54 if self.fallback_value is not None: 

55 return getattr(self.fallback_value, name, None) 

56 raise ImportError( 

57 self.import_error_msg or f"Module {self.module_name} not available", 

58 ) 

59 

60 return getattr(self._module, name) 

61 

62 def _try_import(self) -> None: 

63 """Attempt to import the module.""" 

64 self._import_attempted = True 

65 try: 

66 self._module = importlib.import_module(self.module_name) 

67 _get_logger().debug(f"Successfully imported {self.module_name}") 

68 except ImportError as e: 

69 self._import_failed = True 

70 _get_logger().warning(f"Failed to import {self.module_name}: {e}") 

71 

72 @property 

73 def available(self) -> bool: 

74 """Check if the module is available.""" 

75 if not self._import_attempted: 

76 self._try_import() 

77 return not self._import_failed 

78 

79 def __bool__(self) -> bool: 

80 if not self._import_attempted: 

81 self._try_import() 

82 return not self._import_failed 

83 

84 

85class LazyLoader: 

86 """Manages lazy loading of multiple optional dependencies.""" 

87 

88 def __init__(self) -> None: 

89 self._loaders: dict[str, LazyImport] = {} 

90 

91 def add_import( 

92 self, 

93 name: str, 

94 module_name: str, 

95 fallback_value: Any = None, 

96 error_msg: str | None = None, 

97 ) -> LazyImport: 

98 """Add a lazy import.""" 

99 loader = LazyImport(module_name, fallback_value, error_msg) 

100 self._loaders[name] = loader 

101 return loader 

102 

103 def get_import(self, name: str) -> LazyImport | None: 

104 """Get a lazy import by name.""" 

105 return self._loaders.get(name) 

106 

107 def check_availability(self) -> dict[str, bool]: 

108 """Check availability of all registered imports.""" 

109 return {name: loader.available for name, loader in self._loaders.items()} 

110 

111 

112# Global lazy loader instance 

113lazy_loader = LazyLoader() 

114 

115# Common lazy imports for session-mgmt-mcp 

116transformers = lazy_loader.add_import( 

117 "transformers", 

118 "transformers", 

119 error_msg="Transformers not available. Install with: uv sync --extra embeddings", 

120) 

121 

122onnxruntime = lazy_loader.add_import( 

123 "onnxruntime", 

124 "onnxruntime", 

125 error_msg="ONNX Runtime not available. Install with: uv sync --extra embeddings", 

126) 

127 

128tiktoken = lazy_loader.add_import( 

129 "tiktoken", 

130 "tiktoken", 

131 error_msg="tiktoken not available. Install with: uv sync --extra embeddings", 

132) 

133 

134duckdb = lazy_loader.add_import( 

135 "duckdb", 

136 "duckdb", 

137 error_msg="DuckDB not available. Install with: uv add duckdb", 

138) 

139 

140numpy = lazy_loader.add_import( 

141 "numpy", 

142 "numpy", 

143 error_msg="NumPy not available. Install with: uv sync --extra embeddings", 

144) 

145 

146 

147def require_dependency( 

148 dependency_name: str, 

149 install_hint: str | None = None, 

150) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

151 """Decorator to require a specific dependency for a function.""" 

152 

153 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

154 @wraps(func) 

155 def wrapper(*args: Any, **kwargs: Any) -> Any: 

156 loader = lazy_loader.get_import(dependency_name) 

157 if not loader or not loader.available: 

158 error_msg = f"Function {func.__name__} requires {dependency_name}" 

159 if install_hint: 

160 error_msg += f". Install with: {install_hint}" 

161 raise ImportError(error_msg) 

162 return func(*args, **kwargs) 

163 

164 return wrapper 

165 

166 return decorator 

167 

168 

169def optional_dependency( 

170 dependency_name: str, 

171 fallback_result: Any = None, 

172) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

173 """Decorator to handle optional dependencies gracefully.""" 

174 

175 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

176 @wraps(func) 

177 def wrapper(*args: Any, **kwargs: Any) -> Any: 

178 loader = lazy_loader.get_import(dependency_name) 

179 if not loader or not loader.available: 

180 _get_logger().info( 

181 f"Function {func.__name__} skipped - {dependency_name} not available", 

182 ) 

183 return fallback_result 

184 return func(*args, **kwargs) 

185 

186 return wrapper 

187 

188 return decorator 

189 

190 

191class MockModule: 

192 """Mock module that provides fallback implementations.""" 

193 

194 def __init__(self, name: str) -> None: 

195 self.name = name 

196 

197 def __getattr__(self, name: str) -> Callable[..., Never]: 

198 def mock_function(*args: Any, **kwargs: Any) -> Never: 

199 msg = f"Mock function {name} called - {self.name} not available" 

200 raise ImportError( 

201 msg, 

202 ) 

203 

204 return mock_function 

205 

206 

207def create_embedding_mock() -> Any: 

208 """Create a mock for embedding functionality.""" 

209 

210 class MockEmbedding: 

211 def __init__(self, *args: Any, **kwargs: Any) -> None: 

212 pass 

213 

214 def encode( 

215 self, 

216 texts: str | list[str], 

217 *args: Any, 

218 **kwargs: Any, 

219 ) -> list[list[float]]: 

220 # Return random-like embeddings for testing (using secrets for security) 

221 import secrets 

222 

223 def _random_float() -> float: 

224 """Generate a secure random float between 0 and 1.""" 

225 return secrets.randbelow(1000000) / 1000000.0 

226 

227 if isinstance(texts, str): 

228 return [[_random_float() for _ in range(384)]] 

229 return [[_random_float() for _ in range(384)] for _ in texts] 

230 

231 return MockEmbedding 

232 

233 

234def get_dependency_status() -> dict[str, dict[str, Any]]: 

235 """Get comprehensive status of all dependencies.""" 

236 status: dict[str, dict[str, Any]] = {} 

237 

238 # Check core dependencies 

239 core_deps = ["duckdb"] 

240 for dep in core_deps: 

241 loader = lazy_loader.get_import(dep) 

242 status[dep] = { 

243 "available": loader.available if loader else False, 

244 "required": True, 

245 "category": "core", 

246 } 

247 

248 # Check optional dependencies 

249 optional_deps = ["transformers", "onnxruntime", "tiktoken", "numpy"] 

250 for dep in optional_deps: 

251 loader = lazy_loader.get_import(dep) 

252 status[dep] = { 

253 "available": loader.available if loader else False, 

254 "required": False, 

255 "category": "embeddings" 

256 if dep in {"transformers", "onnxruntime", "numpy"} 

257 else "optimization", 

258 } 

259 

260 # Overall status 

261 core_available = all(status[dep]["available"] for dep in core_deps) 

262 embeddings_available = all( 

263 status[dep]["available"] for dep in ("transformers", "onnxruntime", "numpy") 

264 ) 

265 

266 status["_summary"] = { 

267 "core_functionality": core_available, 

268 "embedding_functionality": embeddings_available, 

269 "optimization_functionality": status["tiktoken"]["available"], 

270 "overall_health": core_available, 

271 } 

272 

273 return status 

274 

275 

276def log_dependency_status() -> None: 

277 """Log the current dependency status.""" 

278 status = get_dependency_status() 

279 summary = status["_summary"] 

280 

281 _get_logger().info( 

282 "Dependency status check completed", 

283 core_functionality=summary["core_functionality"], 

284 embedding_functionality=summary["embedding_functionality"], 

285 optimization_functionality=summary["optimization_functionality"], 

286 ) 

287 

288 # Log missing dependencies 

289 missing_core = [dep for dep in ("duckdb",) if not status[dep]["available"]] 

290 missing_optional = [ 

291 dep 

292 for dep in ("transformers", "onnxruntime", "tiktoken", "numpy") 

293 if not status[dep]["available"] 

294 ] 

295 

296 if missing_core: 

297 _get_logger().warning("Missing core dependencies", missing=missing_core) 

298 

299 if missing_optional: 

300 _get_logger().info( 

301 "Missing optional dependencies", 

302 missing=missing_optional, 

303 install_hint="uv sync --extra embeddings", 

304 ) 

305 

306 

307# Note: Dependency status logging should be called explicitly 

308# to avoid import-time issues. Call log_dependency_status() when needed.