Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-memory/src/lexigram/ai/memory/consolidation/scheduler.py: 23%

1"""Consolidation scheduler — runs periodic consolidation against a MemoryStore.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from typing import TYPE_CHECKING 

7 

8from lexigram.ai.memory.config import ConsolidationConfig 

9from lexigram.contracts.ai.memory import ConsolidationResult 

10from lexigram.logging import ( 

11 get_logger, 

12) 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Sequence 

16 

17 from lexigram.contracts.ai.memory import ( 

18 MemoryConsolidatorProtocol, 

19 MemoryStoreProtocol, 

20 ) 

21 

22logger = get_logger(__name__) 

23 

24 

25class ConsolidationScheduler: 

26 """Runs MemoryConsolidator on a configurable interval. 

27 

28 Designed to be started once and cancelled on shutdown. Consolidation 

29 is only triggered when the interval elapses and entries are available. 

30 The sweep is scoped to an explicit owner set — never all owners. 

31 """ 

32 

33 def __init__( 

34 self, 

35 store: MemoryStoreProtocol, 

36 consolidator: MemoryConsolidatorProtocol, 

37 config: ConsolidationConfig | None = None, 

38 owners: Sequence[str] | None = None, 

39 ) -> None: 

40 """Initialise the scheduler. 

41 

42 Args: 

43 store: Memory store to read entries from. 

44 consolidator: Consolidator run on each cycle. 

45 config: Scheduling configuration. 

46 owners: Explicit owner IDs to consolidate. When ``None`` or 

47 empty the scheduler runs no sweep (no unscoped access). 

48 """ 

49 self._store = store 

50 self._consolidator = consolidator 

51 self._config = config or ConsolidationConfig() 

52 self._owners = owners 

53 self._task: asyncio.Task | None = None 

54 self._background_tasks: set[asyncio.Task] = set() 

55 

56 async def start(self) -> None: 

57 """Start the background consolidation loop.""" 

58 if not self._config.enabled: 

59 logger.info("consolidation_scheduler_disabled") 

60 return 

61 self._task = asyncio.create_task(self._run_loop()) 

62 self._background_tasks.add(self._task) 

63 self._task.add_done_callback(self._background_tasks.discard) 

64 logger.info( 

65 "consolidation_scheduler_started", 

66 interval_s=self._config.interval_seconds, 

67 ) 

68 

69 async def stop(self) -> None: 

70 """Cancel the background consolidation loop.""" 

71 if self._task and not self._task.done(): 

72 self._task.cancel() 

73 try: 

74 await self._task 

75 except asyncio.CancelledError: 

76 pass 

77 self._task = None 

78 logger.info("consolidation_scheduler_stopped") 

79 

80 async def run_once(self) -> ConsolidationResult: 

81 """Execute a single consolidation pass immediately. 

82 

83 Consolidates each owner in the configured owner list; with no 

84 owners configured, returns a zeroed result without touching the 

85 store. 

86 

87 Returns: 

88 Aggregated result of the consolidation passes. 

89 """ 

90 if not self._owners: 

91 logger.info("consolidation_scheduler_no_owners") 

92 return ConsolidationResult( 

93 entries_processed=0, 

94 entries_consolidated=0, 

95 entries_pruned=0, 

96 entities_extracted=0, 

97 duration_ms=0.0, 

98 ) 

99 

100 totals = ConsolidationResult( 

101 entries_processed=0, 

102 entries_consolidated=0, 

103 entries_pruned=0, 

104 entities_extracted=0, 

105 duration_ms=0.0, 

106 ) 

107 processed = consolidated = pruned = extracted = 0 

108 duration = 0.0 

109 for owner_id in self._owners: 

110 entries = await self._store.get_recent( 

111 self._config.batch_size * 10, owner_id 

112 ) 

113 result = await self._consolidator.consolidate(entries) 

114 processed += result.entries_processed 

115 consolidated += result.entries_consolidated 

116 pruned += result.entries_pruned 

117 extracted += result.entities_extracted 

118 duration += result.duration_ms 

119 return ConsolidationResult( 

120 entries_processed=processed, 

121 entries_consolidated=consolidated, 

122 entries_pruned=pruned, 

123 entities_extracted=extracted, 

124 duration_ms=duration, 

125 ) 

126 

127 async def _run_loop(self) -> None: 

128 while True: 

129 await asyncio.sleep(self._config.interval_seconds) 

130 try: 

131 result = await self.run_once() 

132 logger.debug( 

133 "consolidation_cycle", 

134 processed=result.entries_processed, 

135 pruned=result.entries_pruned, 

136 ) 

137 except asyncio.CancelledError: 

138 raise 

139 except MemoryError as exc: 

140 logger.error("consolidation_memory_error", error=str(exc)) 

141 except RuntimeError as exc: 

142 logger.error("consolidation_runtime_error", error=str(exc)) 

143 

144 

145__all__ = ["ConsolidationScheduler"]