Coverage for session_mgmt_mcp/utils/lazy_imports.py: 30.66%
115 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-01 05:22 -0700
1#!/usr/bin/env python3
2"""Lazy loading utilities for optional dependencies.
4This module provides lazy loading for heavy or optional dependencies to improve
5startup performance and handle missing dependencies gracefully.
6"""
8import importlib
9from collections.abc import Callable
10from functools import wraps
11from typing import Any, Never
13from .logging import get_session_logger
15logger = get_session_logger()
18class LazyImport:
19 """Lazy import wrapper that loads modules on first access."""
21 def __init__(
22 self,
23 module_name: str,
24 fallback_value: Any = None,
25 import_error_msg: str | None = None,
26 ) -> None:
27 self.module_name = module_name
28 self.fallback_value = fallback_value
29 self.import_error_msg = import_error_msg
30 self._module = None
31 self._import_attempted = False
32 self._import_failed = False
34 def __getattr__(self, name: str) -> Any:
35 if not self._import_attempted:
36 self._try_import()
38 if self._import_failed:
39 if self.fallback_value is not None:
40 return getattr(self.fallback_value, name, None)
41 raise ImportError(
42 self.import_error_msg or f"Module {self.module_name} not available",
43 )
45 return getattr(self._module, name)
47 def _try_import(self) -> None:
48 """Attempt to import the module."""
49 self._import_attempted = True
50 try:
51 self._module = importlib.import_module(self.module_name)
52 logger.debug(f"Successfully imported {self.module_name}")
53 except ImportError as e:
54 self._import_failed = True
55 logger.warning(f"Failed to import {self.module_name}: {e}")
57 @property
58 def available(self) -> bool:
59 """Check if the module is available."""
60 if not self._import_attempted:
61 self._try_import()
62 return not self._import_failed
64 def __bool__(self) -> bool:
65 if not self._import_attempted:
66 self._try_import()
67 return not self._import_failed
70class LazyLoader:
71 """Manages lazy loading of multiple optional dependencies."""
73 def __init__(self) -> None:
74 self._loaders: dict[str, LazyImport] = {}
76 def add_import(
77 self,
78 name: str,
79 module_name: str,
80 fallback_value: Any = None,
81 error_msg: str | None = None,
82 ) -> LazyImport:
83 """Add a lazy import."""
84 loader = LazyImport(module_name, fallback_value, error_msg)
85 self._loaders[name] = loader
86 return loader
88 def get_import(self, name: str) -> LazyImport | None:
89 """Get a lazy import by name."""
90 return self._loaders.get(name)
92 def check_availability(self) -> dict[str, bool]:
93 """Check availability of all registered imports."""
94 return {name: loader.available for name, loader in self._loaders.items()}
97# Global lazy loader instance
98lazy_loader = LazyLoader()
100# Common lazy imports for session-mgmt-mcp
101transformers = lazy_loader.add_import(
102 "transformers",
103 "transformers",
104 error_msg="Transformers not available. Install with: uv sync --extra embeddings",
105)
107onnxruntime = lazy_loader.add_import(
108 "onnxruntime",
109 "onnxruntime",
110 error_msg="ONNX Runtime not available. Install with: uv sync --extra embeddings",
111)
113tiktoken = lazy_loader.add_import(
114 "tiktoken",
115 "tiktoken",
116 error_msg="tiktoken not available. Install with: uv sync --extra embeddings",
117)
119duckdb = lazy_loader.add_import(
120 "duckdb",
121 "duckdb",
122 error_msg="DuckDB not available. Install with: uv add duckdb",
123)
125numpy = lazy_loader.add_import(
126 "numpy",
127 "numpy",
128 error_msg="NumPy not available. Install with: uv sync --extra embeddings",
129)
132def require_dependency(dependency_name: str, install_hint: str | None = None):
133 """Decorator to require a specific dependency for a function."""
135 def decorator(func: Callable) -> Callable:
136 @wraps(func)
137 def wrapper(*args, **kwargs):
138 loader = lazy_loader.get_import(dependency_name)
139 if not loader or not loader.available:
140 error_msg = f"Function {func.__name__} requires {dependency_name}"
141 if install_hint:
142 error_msg += f". Install with: {install_hint}"
143 raise ImportError(error_msg)
144 return func(*args, **kwargs)
146 return wrapper
148 return decorator
151def optional_dependency(dependency_name: str, fallback_result: Any = None):
152 """Decorator to handle optional dependencies gracefully."""
154 def decorator(func: Callable) -> Callable:
155 @wraps(func)
156 def wrapper(*args, **kwargs):
157 loader = lazy_loader.get_import(dependency_name)
158 if not loader or not loader.available:
159 logger.info(
160 f"Function {func.__name__} skipped - {dependency_name} not available",
161 )
162 return fallback_result
163 return func(*args, **kwargs)
165 return wrapper
167 return decorator
170class MockModule:
171 """Mock module that provides fallback implementations."""
173 def __init__(self, name: str) -> None:
174 self.name = name
176 def __getattr__(self, name: str):
177 def mock_function(*args, **kwargs) -> Never:
178 msg = f"Mock function {name} called - {self.name} not available"
179 raise ImportError(
180 msg,
181 )
183 return mock_function
186def create_embedding_mock():
187 """Create a mock for embedding functionality."""
189 class MockEmbedding:
190 def __init__(self, *args, **kwargs) -> None:
191 pass
193 def encode(self, texts, *args, **kwargs):
194 # Return random-like embeddings for testing
195 import random
197 if isinstance(texts, str):
198 return [[random.random() for _ in range(384)]]
199 return [[random.random() for _ in range(384)] for _ in texts]
201 return MockEmbedding
204def get_dependency_status() -> dict[str, dict[str, Any]]:
205 """Get comprehensive status of all dependencies."""
206 status = {}
208 # Check core dependencies
209 core_deps = ["duckdb"]
210 for dep in core_deps:
211 loader = lazy_loader.get_import(dep)
212 status[dep] = {
213 "available": loader.available if loader else False,
214 "required": True,
215 "category": "core",
216 }
218 # Check optional dependencies
219 optional_deps = ["transformers", "onnxruntime", "tiktoken", "numpy"]
220 for dep in optional_deps:
221 loader = lazy_loader.get_import(dep)
222 status[dep] = {
223 "available": loader.available if loader else False,
224 "required": False,
225 "category": "embeddings"
226 if dep in ["transformers", "onnxruntime", "numpy"]
227 else "optimization",
228 }
230 # Overall status
231 core_available = all(status[dep]["available"] for dep in core_deps)
232 embeddings_available = all(
233 status[dep]["available"] for dep in ["transformers", "onnxruntime", "numpy"]
234 )
236 status["_summary"] = {
237 "core_functionality": core_available,
238 "embedding_functionality": embeddings_available,
239 "optimization_functionality": status["tiktoken"]["available"],
240 "overall_health": core_available,
241 }
243 return status
246def log_dependency_status() -> None:
247 """Log the current dependency status."""
248 status = get_dependency_status()
249 summary = status["_summary"]
251 logger.info(
252 "Dependency status check completed",
253 core_functionality=summary["core_functionality"],
254 embedding_functionality=summary["embedding_functionality"],
255 optimization_functionality=summary["optimization_functionality"],
256 )
258 # Log missing dependencies
259 missing_core = [dep for dep in ["duckdb"] if not status[dep]["available"]]
260 missing_optional = [
261 dep
262 for dep in ["transformers", "onnxruntime", "tiktoken", "numpy"]
263 if not status[dep]["available"]
264 ]
266 if missing_core:
267 logger.warning("Missing core dependencies", missing=missing_core)
269 if missing_optional:
270 logger.info(
271 "Missing optional dependencies",
272 missing=missing_optional,
273 install_hint="uv sync --extra embeddings",
274 )
277# Note: Dependency status logging should be called explicitly
278# to avoid import-time issues. Call log_dependency_status() when needed.