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

1"""DI provider for lexigram-ai-memory.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, cast 

6 

7from lexigram.ai.memory.backends.in_memory import InMemoryMemoryBackend 

8from lexigram.ai.memory.config import MemoryConfig 

9from lexigram.ai.memory.consolidation.consolidator import MemoryConsolidator 

10from lexigram.ai.memory.consolidation.scheduler import ConsolidationScheduler 

11from lexigram.ai.memory.episodic.store import EpisodicMemoryStore 

12from lexigram.ai.memory.pruning.pruner import DynamicContextPruner 

13from lexigram.ai.memory.retrieval.ranking import RelevanceRanker 

14from lexigram.ai.memory.retrieval.retriever import MemoryRetriever 

15from lexigram.ai.memory.semantic.store import SemanticMemoryStore 

16from lexigram.ai.memory.working.manager import WorkingMemoryManager 

17from lexigram.contracts.ai.memory import ( 

18 EpisodicMemoryProtocol, 

19 MemoryConsolidatorProtocol, 

20 MemoryStoreProtocol, 

21 SemanticMemoryProtocol, 

22 WorkingMemoryProtocol, 

23) 

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

25from lexigram.di.provider import Provider, ProviderPriority 

26from lexigram.logging import ( 

27 get_logger, 

28) 

29 

30if TYPE_CHECKING: 

31 from lexigram.contracts.core.di import ( 

32 ContainerRegistrarProtocol, 

33 ContainerResolverProtocol, 

34 ) 

35 

36logger = get_logger(__name__) 

37 

38 

39class MemoryProvider(Provider): 

40 """Registers all memory services in the DI container. 

41 

42 Wires up the MemoryStoreProtocol, EpisodicMemoryProtocol, 

43 SemanticMemoryProtocol, WorkingMemoryProtocol, and consolidation 

44 services. The default backend is ``InMemoryMemoryBackend``; callers 

45 with persistent needs should register a ``MemoryStoreProtocol`` 

46 override after this provider runs. 

47 

48 Note: 

49 All memory services are owner-scoped — entries carry an 

50 ``owner_id`` and every store operation is restricted to one 

51 owner (user, session, or system scope). 

52 """ 

53 

54 name = "ai-memory" 

55 priority = ProviderPriority.DOMAIN 

56 config_key: str | None = "ai_memory" 

57 config_model: type | None = MemoryConfig 

58 

59 def __init__( 

60 self, 

61 config: MemoryConfig | None = None, 

62 enable_consolidation: bool = True, 

63 ) -> None: 

64 """Initialise the provider. 

65 

66 Args: 

67 config: Memory configuration. Defaults to ``MemoryConfig()``. 

68 enable_consolidation: Start the consolidation scheduler during 

69 :meth:`boot`. When ``False``, the scheduler is never started 

70 regardless of ``config.consolidation.enabled``. 

71 """ 

72 super().__init__() 

73 self._requested_config = config 

74 self._mem_config = config or MemoryConfig() 

75 self._enable_consolidation = enable_consolidation 

76 self._scheduler: ConsolidationScheduler | None = None 

77 

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

79 """Register memory services. 

80 

81 Args: 

82 container: DI container registrar. 

83 """ 

84 self._mem_config = self._requested_config or ( 

85 self.config 

86 if isinstance(getattr(self, "config", None), MemoryConfig) 

87 else self._mem_config 

88 ) 

89 container.singleton(MemoryConfig, self._mem_config) 

90 

91 if not self._mem_config.enabled: 

92 logger.info("memory_disabled", reason="MemoryConfig.enabled=False") 

93 return 

94 

95 # Follow-up: wire cache/database/vector backends at boot phase 

96 # (see 2026-08-12-ai-family-gaps-remediation.md Task 5) — requires 

97 # resolving CacheBackendProtocol/DatabaseProviderProtocol from the 

98 # container, which ContainerRegistrarProtocol does not permit. 

99 if self._mem_config.default_backend != "in_memory": 

100 logger.warning( 

101 "memory_backend_not_implemented", 

102 requested_backend=self._mem_config.default_backend, 

103 fallback="in_memory", 

104 ) 

105 

106 backend = InMemoryMemoryBackend() 

107 episodic = EpisodicMemoryStore(backend=backend) 

108 semantic = SemanticMemoryStore( 

109 min_confidence=self._mem_config.semantic.min_confidence, 

110 max_facts_per_entity=self._mem_config.semantic.max_facts_per_entity, 

111 ) 

112 working = WorkingMemoryManager( 

113 episodic=episodic, 

114 semantic=cast("SemanticMemoryProtocol", semantic), 

115 config=self._mem_config.working, 

116 ) 

117 consolidator = MemoryConsolidator(config=self._mem_config.consolidation) 

118 retriever = MemoryRetriever(sources=[backend]) 

119 

120 container.singleton(MemoryStoreProtocol, lambda: backend) 

121 container.singleton(EpisodicMemoryProtocol, lambda: episodic) 

122 container.singleton(SemanticMemoryProtocol, lambda: semantic) 

123 container.singleton(WorkingMemoryProtocol, lambda: working) 

124 container.singleton(MemoryConsolidatorProtocol, lambda: consolidator) 

125 container.singleton(MemoryRetriever, lambda: retriever) 

126 container.singleton(RelevanceRanker, RelevanceRanker) 

127 # Register pruner — the container will auto-inject TokenCounterProtocol 

128 container.singleton(DynamicContextPruner, DynamicContextPruner) 

129 

130 logger.debug("memory_services_registered") 

131 

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

133 """Start the consolidation scheduler if enabled. 

134 

135 Args: 

136 container: Resolved DI container. 

137 """ 

138 if not self._mem_config.enabled: 

139 return 

140 

141 if self._enable_consolidation and self._mem_config.consolidation.enabled: 

142 store = await container.resolve_optional(MemoryStoreProtocol) 

143 consolidator = await container.resolve_optional(MemoryConsolidatorProtocol) 

144 if store is None or consolidator is None: 

145 logger.warning("memory_consolidation_dependencies_missing") 

146 return 

147 

148 self._scheduler = ConsolidationScheduler( 

149 store=store, 

150 consolidator=consolidator, 

151 config=self._mem_config.consolidation, 

152 ) 

153 await self._scheduler.start() 

154 logger.info("memory_consolidation_scheduler_started") 

155 

156 async def shutdown(self) -> None: 

157 """Stop the consolidation scheduler.""" 

158 if self._scheduler: 

159 await self._scheduler.stop() 

160 self._scheduler = None 

161 

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

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

164 

165 No external backend to ping. 

166 

167 Args: 

168 timeout: Ignored for in-process providers. 

169 

170 Returns: 

171 Always HEALTHY — no external backend to ping. 

172 """ 

173 return HealthCheckResult( 

174 component=self.name, 

175 status=HealthStatus.HEALTHY, 

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

177 ) 

178 

179 

180__all__ = ["MemoryProvider"]