Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-workers/src/lexigram/ai/workers/di/provider.py: 52%

1"""Workers DI provider — registers background workers with the container.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from typing import TYPE_CHECKING, Any 

7 

8from lexigram.ai.workers.batch_embedding.worker import BatchEmbeddingWorker 

9from lexigram.ai.workers.config import WorkersConfig 

10from lexigram.ai.workers.dlq.worker import DeadLetterQueueWorker 

11from lexigram.ai.workers.document_ingestion.worker import DocumentIngestionWorker 

12from lexigram.ai.workers.maintenance.worker import MaintenanceWorker 

13from lexigram.contracts.core.health import HealthCheckResult, HealthStatus 

14from lexigram.contracts.core.provider import ProviderPriority 

15from lexigram.contracts.exceptions.container import UnresolvableDependencyError 

16from lexigram.di.provider import Provider 

17from lexigram.logging import ( 

18 get_logger, 

19) 

20 

21if TYPE_CHECKING: 

22 from lexigram.contracts.core.di import ( 

23 ContainerRegistrarProtocol, 

24 ContainerResolverProtocol, 

25 ) 

26 

27logger = get_logger(__name__) 

28 

29# Worker types managed by this provider 

30_WORKER_TYPES: tuple[type, ...] = ( 

31 BatchEmbeddingWorker, 

32 DocumentIngestionWorker, 

33 MaintenanceWorker, 

34 DeadLetterQueueWorker, 

35) 

36 

37 

38class WorkersProvider(Provider): 

39 """Provider for AI background workers. 

40 

41 Reads :class:`~lexigram.ai.workers.config.WorkersConfig` and registers 

42 workers for batch embedding, ingestion, maintenance, and DLQ handling. 

43 

44 Lifecycle: 

45 - ``register`` — binds worker singletons into the container. 

46 - ``boot`` — resolves workers and starts each as a background async task. 

47 - ``shutdown`` — calls ``stop()`` on every worker, then awaits each 

48 background task to ensure graceful queue draining. 

49 """ 

50 

51 name = "workers" 

52 priority = ProviderPriority.INFRASTRUCTURE 

53 config_key: str | None = "ai_workers" 

54 config_model: type | None = WorkersConfig 

55 

56 def __init__( 

57 self, 

58 config: WorkersConfig | None = None, 

59 enable_scheduler: bool = True, 

60 ) -> None: 

61 super().__init__() 

62 self._requested_config = config 

63 self._config = config or WorkersConfig() 

64 self._enable_scheduler = enable_scheduler 

65 self._workers: list[Any] = [] 

66 self._tasks: set[asyncio.Task[None]] = set() 

67 

68 @classmethod 

69 def from_config(cls, config: WorkersConfig, **context) -> WorkersProvider: 

70 """Factory method for DI container setup.""" 

71 return cls(config) 

72 

73 async def register(self, container: ContainerRegistrarProtocol) -> None: 

74 """Register workers with the DI container.""" 

75 self._config = self._requested_config or ( 

76 self.config 

77 if isinstance(getattr(self, "config", None), WorkersConfig) 

78 else self._config 

79 ) 

80 container.singleton(WorkersConfig, self._config) 

81 

82 if not self._config.enabled: 

83 logger.info( 

84 "workers_provider_disabled", reason="WorkersConfig.enabled=False" 

85 ) 

86 return 

87 

88 for worker_type in _WORKER_TYPES: 

89 if hasattr(container, "register"): 

90 container.register(worker_type) 

91 else: 

92 container.transient(worker_type, worker_type) 

93 

94 logger.info("workers_provider_registered", worker_count=len(_WORKER_TYPES)) 

95 

96 async def boot(self, container: ContainerResolverProtocol) -> None: 

97 """Resolve workers from the container and start each as a background task. 

98 

99 Each worker's ``start()`` coroutine is wrapped in an ``asyncio.Task`` 

100 so it runs concurrently. Task references are stored in ``_tasks`` to 

101 prevent garbage collection (Ruff RUF006). 

102 

103 Args: 

104 container: Container resolver for resolving worker instances. 

105 """ 

106 if not self._config.enabled: 

107 logger.debug("workers_provider_boot_skipped", reason="disabled") 

108 return 

109 

110 for worker_type in _WORKER_TYPES: 

111 try: 

112 worker: Any = await container.resolve(worker_type) 

113 except ( 

114 LookupError, 

115 RuntimeError, 

116 TypeError, 

117 ValueError, 

118 AttributeError, 

119 UnresolvableDependencyError, 

120 ) as exc: 

121 logger.warning( 

122 "worker_resolve_failed", 

123 worker=worker_type.__name__, 

124 error=str(exc), 

125 ) 

126 continue 

127 

128 self._workers.append(worker) 

129 task = asyncio.create_task( 

130 worker.start(), 

131 name=f"worker.{worker_type.__name__}", 

132 ) 

133 self._tasks.add(task) 

134 task.add_done_callback(self._on_worker_done) 

135 logger.info("worker_started", worker=worker_type.__name__) 

136 

137 logger.info("workers_provider_booted", workers_started=len(self._workers)) 

138 

139 async def shutdown(self) -> None: 

140 """Stop all workers gracefully and await their background tasks. 

141 

142 Calls ``stop()`` on every worker in reverse start order, then waits 

143 for each background task to finish (with a 30 s timeout per task). 

144 """ 

145 for worker in reversed(self._workers): 

146 try: 

147 await worker.stop() 

148 except (OSError, RuntimeError) as exc: 

149 logger.warning( 

150 "worker_stop_error", 

151 worker=type(worker).__name__, 

152 error=str(exc), 

153 ) 

154 

155 if self._tasks: 

156 await asyncio.wait(self._tasks, timeout=30.0) 

157 pending = {t for t in self._tasks if not t.done()} 

158 for task in pending: 

159 task.cancel() 

160 logger.warning("worker_task_cancelled", task=task.get_name()) 

161 

162 self._workers.clear() 

163 self._tasks.clear() 

164 logger.info("workers_provider_shutdown") 

165 

166 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

167 """Health check — always healthy (in-process domain provider). 

168 

169 No external backend to ping. 

170 

171 Args: 

172 timeout: Ignored for in-process providers. 

173 

174 Returns: 

175 Always HEALTHY — no external backend to ping. 

176 """ 

177 return HealthCheckResult( 

178 component=self.name, 

179 status=HealthStatus.HEALTHY, 

180 details={"status": "operational"}, 

181 ) 

182 

183 # ------------------------------------------------------------------ 

184 # Internal helpers 

185 # ------------------------------------------------------------------ 

186 

187 def _on_worker_done(self, task: asyncio.Task[None]) -> None: 

188 """Callback to handle completed or failed worker tasks.""" 

189 self._tasks.discard(task) 

190 if task.cancelled(): 

191 logger.info("worker_task_cancelled", task=task.get_name()) 

192 elif exc := task.exception(): 

193 logger.error( 

194 "worker_task_failed", 

195 task=task.get_name(), 

196 error=str(exc), 

197 ) 

198 

199 

200__all__ = ["WorkersProvider"]