1#!/usr/bin/env python3

2"""Lazy loading utilities for optional dependencies.

3

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

5 startup performance and handle missing dependencies gracefully.

6"""

7

8 import importlib

9 from collections.abc import Callable

10 from functools import wraps

11 from typing import Any, Never

12

13 from .logging import get_session_logger

14

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

16_logger: Any = None

17

18

19 def _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

33 class 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

85 class 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

113 lazy_loader = LazyLoader()

114

115# Common lazy imports for session-mgmt-mcp

116 transformers = lazy_loader.add_import(

117 "transformers",

118 "transformers",

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

120)

121

122 onnxruntime = lazy_loader.add_import(

123 "onnxruntime",

124 "onnxruntime",

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

126)

127

128 tiktoken = lazy_loader.add_import(

129 "tiktoken",

130 "tiktoken",

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

132)

133

134 duckdb = lazy_loader.add_import(

135 "duckdb",

136 "duckdb",

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

138)

139

140 numpy = lazy_loader.add_import(

141 "numpy",

142 "numpy",

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

144)

145

146

147 def 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

169 def 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

191 class 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

207 def 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

234 def 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

276 def 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.