« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""
2Maintenance worker for periodic cleanup and optimization.
4Handles scheduled maintenance tasks including:
5- Vector store index optimization
6- Embedding cache cleanup (TTL-based)
7- Old document cleanup
8- Metrics aggregation and rollup
9- Health check monitoring
10"""
12from __future__ import annotations
14import asyncio
15import contextlib
16from datetime import UTC, datetime
17from typing import TYPE_CHECKING, Any
19from lexigram.ai.workers.types import (
20 MaintenanceResult,
21 MaintenanceStatus,
22 MaintenanceTask,
23 MaintenanceTaskType,
24)
25from lexigram.contracts.core.health import HealthCheckResult, HealthStatus
26from lexigram.logging import (
27 get_logger,
28)
30if TYPE_CHECKING:
31 from collections.abc import Callable
33 from lexigram.contracts import VectorStoreProtocol
35logger = get_logger(__name__)
38class MaintenanceWorker:
39 """
40 Background worker for scheduled maintenance tasks.
42 Executes periodic maintenance operations including:
43 - Vector store index optimization
44 - Cache cleanup and TTL enforcement
45 - Old document cleanup
46 - Metrics aggregation
48 Example:
49 ```python
50 from lexigram.ai.workers import MaintenanceWorker
51 from lexigram.contracts import VectorStoreProtocol
53 # Setup — resolve via the DI container
54 vector_store = container.resolve(VectorStoreProtocol)
56 worker = MaintenanceWorker(
57 vector_store=vector_store,
58 worker_id="maintenance",
59 )
61 # Register maintenance tasks
62 worker.register_task(
63 name="optimize_indexes",
64 task_type=MaintenanceTaskType.INDEX_OPTIMIZATION,
65 handler=lambda: vector_store.optimize_indexes(),
66 interval_seconds=3600, # Every hour
67 )
69 worker.register_task(
70 name="cleanup_old_docs",
71 task_type=MaintenanceTaskType.DOCUMENT_CLEANUP,
72 handler=lambda: cleanup_documents(days=30),
73 schedule_cron="0 2 * * *", # 2 AM daily
74 )
76 # Start worker
77 await worker.start()
79 # Get statistics
80 stats = worker.get_stats()
81 logger.info(f"Tasks run: {stats['total_runs']}")
83 # Stop worker
84 await worker.stop()
85 ```
86 """
88 def __init__(
89 self,
90 vector_store: VectorStoreProtocol | None = None,
91 worker_id: str = "maintenance",
92 check_interval: int = 60, # Check every 60 seconds
93 ):
94 """
95 Initialize maintenance worker.
97 Args:
98 vector_store: Optional vector store for index optimization
99 worker_id: Unique worker identifier
100 check_interval: Interval in seconds to check for tasks to run
101 """
102 self.vector_store = vector_store
103 self.worker_id = worker_id
104 self.check_interval = check_interval
106 # Task registry
107 self._tasks: dict[str, MaintenanceTask] = {}
108 self._tasks_lock = asyncio.Lock()
110 # Execution history
111 self._history: list[MaintenanceResult] = []
112 self._history_lock = asyncio.Lock()
113 self._max_history = 1000 # Keep last 1000 results
115 # Worker state
116 self._running = False
117 self._worker_task: asyncio.Task | None = None
119 async def start(self) -> None:
120 """Start the maintenance worker."""
121 if self._running:
122 logger.warning("Worker %s already running", self.worker_id)
123 return
125 self._running = True
126 self._worker_task = asyncio.create_task(self._maintenance_loop())
128 logger.info(
129 "Started maintenance worker",
130 worker_id=self.worker_id,
131 check_interval=self.check_interval,
132 registered_tasks=len(self._tasks),
133 )
135 async def stop(self) -> None:
136 """Stop the maintenance worker."""
137 if not self._running:
138 return
140 self._running = False
142 if self._worker_task:
143 self._worker_task.cancel()
144 with contextlib.suppress(asyncio.CancelledError):
145 await self._worker_task
147 logger.info("Stopped maintenance worker", worker_id=self.worker_id)
149 def register_task(
150 self,
151 name: str,
152 task_type: MaintenanceTaskType,
153 handler: Callable[[], Any],
154 schedule_cron: str | None = None,
155 interval_seconds: int | None = None,
156 enabled: bool = True,
157 timeout: float = 300.0,
158 ) -> None:
159 """
160 Register a maintenance task.
162 Args:
163 name: Unique task name
164 task_type: Type of maintenance task
165 handler: Async or sync callable to execute
166 schedule_cron: Cron expression for scheduling
167 interval_seconds: Simple interval in seconds
168 enabled: Whether task is enabled
169 timeout: Task timeout in seconds
170 """
171 if not schedule_cron and not interval_seconds:
172 msg = "Must provide either schedule_cron or interval_seconds"
173 raise ValueError(msg)
175 task = MaintenanceTask(
176 name=name,
177 task_type=task_type,
178 handler=handler,
179 schedule_cron=schedule_cron,
180 interval_seconds=interval_seconds,
181 enabled=enabled,
182 timeout=timeout,
183 )
185 self._tasks[name] = task
187 logger.info(
188 "Registered maintenance task",
189 task_name=name,
190 task_type=task_type.value,
191 interval_seconds=interval_seconds,
192 cron=schedule_cron,
193 )
195 def unregister_task(self, name: str) -> None:
196 """Unregister a maintenance task."""
197 if name in self._tasks:
198 del self._tasks[name]
199 logger.info("Unregistered maintenance task", task_name=name)
201 def enable_task(self, name: str) -> None:
202 """Enable a maintenance task."""
203 if name in self._tasks:
204 self._tasks[name].enabled = True
205 logger.info("Enabled maintenance task", task_name=name)
207 def disable_task(self, name: str) -> None:
208 """Disable a maintenance task."""
209 if name in self._tasks:
210 self._tasks[name].enabled = False
211 logger.info("Disabled maintenance task", task_name=name)
213 async def run_task_now(self, name: str) -> MaintenanceResult:
214 """
215 Run a specific maintenance task immediately.
217 Args:
218 name: Task name to run
220 Returns:
221 Maintenance result
222 """
223 async with self._tasks_lock:
224 if name not in self._tasks:
225 msg = f"Task not found: {name}"
226 raise ValueError(msg)
228 task = self._tasks[name]
230 logger.info("Running maintenance task manually", task_name=name)
231 return await self._execute_task(task)
233 async def _maintenance_loop(self) -> None:
234 """Main maintenance loop that checks and runs tasks."""
235 while self._running:
236 try:
237 # Check all tasks for execution
238 async with self._tasks_lock:
239 tasks_to_run = [
240 task for task in self._tasks.values() if task.should_run()
241 ]
243 # Execute tasks that need to run
244 for task in tasks_to_run:
245 try:
246 result = await self._execute_task(task)
248 # Store result
249 await self._store_result(result)
251 # Update task
252 async with self._tasks_lock:
253 task.last_run = result.completed_at
254 task.last_status = result.status
255 task.last_error = result.error
256 task.run_count += 1
258 except Exception as e:
259 logger.exception(
260 "Error executing maintenance task",
261 task_name=task.name,
262 error=str(e),
263 )
265 # Wait for next check
266 await asyncio.sleep(self.check_interval)
268 except asyncio.CancelledError:
269 break
270 except Exception as e:
271 logger.exception(
272 "Error in maintenance loop",
273 error=str(e),
274 )
275 await asyncio.sleep(self.check_interval)
277 async def _execute_task(self, task: MaintenanceTask) -> MaintenanceResult:
278 """Execute a single maintenance task."""
279 started_at = datetime.now(UTC)
281 logger.info(
282 "Executing maintenance task",
283 task_name=task.name,
284 task_type=task.task_type.value,
285 )
287 try:
288 # Execute with timeout
289 if asyncio.iscoroutinefunction(task.handler):
290 result = await asyncio.wait_for(
291 task.handler(),
292 timeout=task.timeout,
293 )
294 else:
295 # Run sync handler in thread pool
296 result = await asyncio.wait_for(
297 asyncio.to_thread(task.handler),
298 timeout=task.timeout,
299 )
301 # Parse result
302 items_processed = 0
303 items_deleted = 0
304 metadata = {}
306 if isinstance(result, dict):
307 items_processed = result.get("items_processed", 0)
308 items_deleted = result.get("items_deleted", 0)
309 metadata = result.get("metadata", {})
310 elif isinstance(result, int):
311 items_processed = result
313 maintenance_result = MaintenanceResult.success(
314 task_name=task.name,
315 task_type=task.task_type,
316 started_at=started_at,
317 items_processed=items_processed,
318 items_deleted=items_deleted,
319 metadata=metadata,
320 )
322 logger.info(
323 "Maintenance task completed",
324 task_name=task.name,
325 duration=f"{maintenance_result.duration_seconds:.2f}s",
326 items_processed=items_processed,
327 items_deleted=items_deleted,
328 )
330 except TimeoutError:
331 error_msg = f"Task timed out after {task.timeout}s"
332 logger.exception(
333 "Maintenance task timed out",
334 timeout=task.timeout,
335 )
337 return MaintenanceResult.failure(
338 task_name=task.name,
339 task_type=task.task_type,
340 started_at=started_at,
341 error=error_msg,
342 )
344 except Exception as e:
345 error_msg = str(e)
346 logger.exception(
347 "Maintenance task failed",
348 task_name=task.name,
349 error=error_msg,
350 )
352 return MaintenanceResult.failure(
353 task_name=task.name,
354 task_type=task.task_type,
355 started_at=started_at,
356 error=error_msg,
357 )
358 else:
359 return maintenance_result
361 async def _store_result(self, result: MaintenanceResult) -> None:
362 """Store maintenance result in history."""
363 async with self._history_lock:
364 self._history.append(result)
366 # Trim history if needed
367 if len(self._history) > self._max_history:
368 self._history = self._history[-self._max_history :]
370 def get_stats(self) -> dict[str, Any]:
371 """Get worker statistics."""
372 total_runs = sum(task.run_count for task in self._tasks.values())
373 successful_runs = sum(
374 1
375 for result in self._history
376 if result.status == MaintenanceStatus.COMPLETED
377 )
378 failed_runs = sum(
379 1 for result in self._history if result.status == MaintenanceStatus.FAILED
380 )
382 return {
383 "worker_id": self.worker_id,
384 "running": self._running,
385 "registered_tasks": len(self._tasks),
386 "total_runs": total_runs,
387 "successful_runs": successful_runs,
388 "failed_runs": failed_runs,
389 "history_size": len(self._history),
390 }
392 def get_task_status(self, name: str) -> dict[str, Any] | None:
393 """Get status of a specific task."""
394 if name not in self._tasks:
395 return None
397 return self._tasks[name].to_dict()
399 def get_all_tasks(self) -> list[dict[str, Any]]:
400 """Get status of all registered tasks."""
401 return [task.to_dict() for task in self._tasks.values()]
403 def get_recent_results(self, limit: int = 10) -> list[dict[str, Any]]:
404 """Get recent maintenance results."""
405 return [result.to_dict() for result in self._history[-limit:]]
407 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
408 """Report the health of this worker.
410 Args:
411 timeout: Unused; present for protocol conformance.
413 Returns:
414 HEALTHY when the worker is running, UNHEALTHY otherwise.
415 """
416 status = HealthStatus.HEALTHY if self._running else HealthStatus.UNHEALTHY
417 stats = self.get_stats()
418 return HealthCheckResult(
419 component=f"worker.maintenance.{self.worker_id}",
420 status=status,
421 details=stats,
422 )
424 # Built-in maintenance handlers
426 async def optimize_vector_indexes(self) -> dict[str, Any]:
427 """
428 Optimize vector store indexes.
430 This is a built-in handler for index optimization.
431 """
432 if not self.vector_store:
433 msg = "Vector store not configured"
434 raise ValueError(msg)
436 logger.info("Starting vector index optimization")
438 # Call vector store optimization if supported
439 if hasattr(self.vector_store, "optimize_indexes"):
440 result = await self.vector_store.optimize_indexes()
441 return {
442 "items_processed": result.get("collections_optimized", 0),
443 "metadata": result,
444 }
446 logger.warning("Vector store does not support index optimization")
447 return {"items_processed": 0}
449 async def cleanup_old_embeddings_cache(
450 self,
451 max_age_days: int = 30,
452 ) -> dict[str, Any]:
453 """
454 Clean up old embeddings from cache.
456 This is a built-in handler for cache cleanup.
458 Args:
459 max_age_days: Maximum age in days for cached embeddings
460 """
461 # This is a placeholder - actual implementation would need
462 # access to the embedding cache
463 logger.info(
464 "Cleaning up embeddings cache",
465 max_age_days=max_age_days,
466 )
468 # Placeholder return
469 return {
470 "items_deleted": 0,
471 "metadata": {"max_age_days": max_age_days},
472 }
474 async def aggregate_metrics(self) -> dict[str, Any]:
475 """
476 Aggregate and rollup metrics.
478 This is a built-in handler for metrics aggregation.
479 """
480 logger.info("Aggregating metrics")
482 # Placeholder - actual implementation would aggregate
483 # metrics from the metrics collector
484 return {
485 "items_processed": 0,
486 "metadata": {"aggregation_type": "hourly"},
487 }