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

1"""DI Provider for the RAG (Retrieval Augmented Generation) subsystem.""" 

2 

3from __future__ import annotations 

4 

5import importlib.util 

6from typing import TYPE_CHECKING, Any 

7 

8from lexigram.ai.rag.config import RAGConfig 

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

10from lexigram.contracts.exceptions.container import UnresolvableDependencyError 

11from lexigram.contracts.exceptions.provider import ModuleVisibilityError 

12from lexigram.di.provider import Provider, ProviderPriority 

13from lexigram.logging import ( 

14 get_logger, 

15) 

16 

17if TYPE_CHECKING: 

18 from lexigram.contracts.core.di import ( 

19 BootContainerProtocol, 

20 ContainerRegistrarProtocol, 

21 ) 

22 

23logger = get_logger(__name__) 

24 

25 

26def _llmlingua_available() -> bool: 

27 """Check if llmlingua package is installed.""" 

28 try: 

29 return importlib.util.find_spec("llmlingua") is not None 

30 except (ValueError, AttributeError): 

31 return False 

32 

33 

34def _flashrank_available() -> bool: 

35 """Check if flashrank package is installed.""" 

36 try: 

37 return importlib.util.find_spec("flashrank") is not None 

38 except (ValueError, AttributeError): 

39 return False 

40 

41 

42class RAGProvider(Provider): 

43 """Registers RAG pipeline services and strategy registries with the DI container.""" 

44 

45 name = "rag" 

46 priority = ProviderPriority.DOMAIN 

47 config_key: str | None = "ai_rag" 

48 config_model: type | None = RAGConfig 

49 

50 def __init__(self, config: RAGConfig | None = None) -> None: 

51 super().__init__() 

52 self._requested_config = config 

53 self._config = config or RAGConfig() 

54 

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

56 self._config = self._requested_config or self._config or RAGConfig() 

57 container.singleton(RAGConfig, instance=self._config) 

58 

59 from lexigram.ai.rag.context_compression.strategy_registry import ( 

60 CompressionStrategyRegistry, 

61 ) 

62 

63 compression_registry = CompressionStrategyRegistry.with_defaults() 

64 

65 # Register LLMLingua-2 handler if available 

66 if _llmlingua_available(): 

67 from lexigram.ai.rag.context_compression.strategies.llmlingua2 import ( 

68 LLMLingua2CompressorStrategy, 

69 LLMLingua2StrategyHandler, 

70 ) 

71 

72 compression_registry.register( 

73 LLMLingua2StrategyHandler(LLMLingua2CompressorStrategy()) 

74 ) 

75 logger.debug("llmlingua2_compressor_registered") 

76 else: 

77 logger.debug("llmlingua2_compressor_skipped_not_installed") 

78 

79 container.singleton(CompressionStrategyRegistry, instance=compression_registry) 

80 

81 from lexigram.ai.rag.hyde.strategy_registry import HyDEStrategyRegistry 

82 

83 hyde_registry = HyDEStrategyRegistry.with_defaults() 

84 container.singleton(HyDEStrategyRegistry, instance=hyde_registry) 

85 

86 from lexigram.ai.rag.reasoning.strategy_registry import ( 

87 ReasoningStrategyRegistry, 

88 ) 

89 

90 reasoning_registry = ReasoningStrategyRegistry.with_defaults() 

91 container.singleton(ReasoningStrategyRegistry, instance=reasoning_registry) 

92 

93 from lexigram.ai.rag.pipeline.stages.synthesis_registry import ( 

94 SynthesisStrategyRegistry, 

95 ) 

96 

97 synthesis_registry = SynthesisStrategyRegistry.with_defaults() 

98 container.singleton(SynthesisStrategyRegistry, instance=synthesis_registry) 

99 

100 from lexigram.ai.rag.chunking.strategy_registry import ( 

101 ChunkingStrategyRegistry, 

102 ) 

103 

104 chunking_registry = ChunkingStrategyRegistry.with_defaults() 

105 container.singleton(ChunkingStrategyRegistry, instance=chunking_registry) 

106 

107 from lexigram.ai.rag.reranking.strategy_registry import ( 

108 RerankingStrategyRegistry, 

109 ) 

110 

111 reranking_registry = RerankingStrategyRegistry.with_defaults() 

112 

113 # Register FlashRank handler if available 

114 if _flashrank_available(): 

115 from lexigram.ai.rag.reranking.strategies.flashrank import ( 

116 FlashRankStrategyHandler, 

117 ) 

118 

119 reranking_registry.register(FlashRankStrategyHandler()) 

120 logger.debug("flashrank_reranker_registered") 

121 else: 

122 logger.debug("flashrank_reranker_skipped_not_installed") 

123 

124 container.singleton(RerankingStrategyRegistry, instance=reranking_registry) 

125 

126 await self._discover_strategies(container) 

127 

128 await self._register_knowledge_graph(container) 

129 

130 logger.info("rag_provider_registered") 

131 

132 async def _register_knowledge_graph( 

133 self, 

134 container: ContainerRegistrarProtocol, 

135 ) -> None: 

136 """Register the default in-memory knowledge graph singleton.""" 

137 from lexigram.ai.rag.knowledge_graph.core import KnowledgeGraph 

138 

139 container.singleton(KnowledgeGraph, instance=KnowledgeGraph()) 

140 logger.info("rag_knowledge_graph_in_memory") 

141 

142 async def _discover_strategies(self, container: ContainerRegistrarProtocol) -> None: 

143 """Auto-discover RAG strategy providers via entry-points.""" 

144 import importlib.metadata as _meta 

145 

146 from lexigram.di.provider import Provider as _Provider 

147 

148 for group in ( 

149 "lexigram.chunking.strategies", 

150 "lexigram.retrieval.strategies", 

151 ): 

152 eps = _meta.entry_points(group=group) 

153 for ep in eps: 

154 try: 

155 candidate = ep.load() 

156 except (ImportError, AttributeError, TypeError, ValueError) as exc: 

157 logger.warning( 

158 "rag_strategy_ep_load_failed", 

159 name=ep.name, 

160 group=group, 

161 error=str(exc), 

162 ) 

163 continue 

164 if isinstance(candidate, type) and issubclass(candidate, _Provider): 

165 await candidate().register(container) 

166 logger.info("rag_strategy_ep_loaded", name=ep.name, group=group) 

167 else: 

168 logger.debug("rag_strategy_ep_skipped", name=ep.name, group=group) 

169 

170 async def _maybe_wrap_with_tenancy( 

171 self, 

172 container: BootContainerProtocol, 

173 ) -> None: 

174 """Register ``TenantScopedRAGPipeline`` when tenancy is enabled.""" 

175 if not self._config.tenancy.enabled: 

176 return 

177 

178 from lexigram.ai.rag.tenancy import TenantScopedRAGPipeline 

179 from lexigram.ai.rag.tenancy.resolver import ( 

180 TemplatedTenantCollectionResolver, 

181 ) 

182 from lexigram.contracts.ai.rag import RAGPipelineProtocol 

183 from lexigram.primitives.context import Context 

184 

185 ctx = await container.resolve(Context) 

186 resolver: Any = TemplatedTenantCollectionResolver() 

187 base_config = self._config 

188 

189 async def _default_factory(config: RAGConfig) -> RAGPipelineProtocol: 

190 from lexigram.ai.rag.config import PipelineConfig 

191 from lexigram.ai.rag.pipeline import RAGPipeline 

192 from lexigram.ai.rag.pipeline._stage_factory import ( 

193 build_pipeline_stages, 

194 ) 

195 

196 pipeline_cfg = PipelineConfig() 

197 stages = build_pipeline_stages(pipeline_cfg) 

198 return RAGPipeline(config=pipeline_cfg, stages=stages) 

199 

200 container.singleton( 

201 RAGPipelineProtocol, 

202 instance=TenantScopedRAGPipeline( 

203 base_config=base_config, 

204 resolver=resolver, 

205 ctx=ctx, 

206 pipeline_factory=_default_factory, 

207 ), 

208 ) 

209 logger.info("rag_tenancy_enabled") 

210 

211 async def boot(self, container: BootContainerProtocol) -> None: 

212 """Boot RAG provider — wire optional integrations.""" 

213 self._booted_container = container 

214 

215 # Optional: tenancy 

216 await self._maybe_wrap_with_tenancy(container) 

217 

218 # Optional: working memory (from lexigram-ai-memory) 

219 working_memory = None 

220 try: 

221 from lexigram.contracts.ai.memory import WorkingMemoryProtocol 

222 

223 working_memory = await container.resolve(WorkingMemoryProtocol) 

224 logger.debug("rag_working_memory_available") 

225 except ( 

226 LookupError, 

227 RuntimeError, 

228 AttributeError, 

229 ImportError, 

230 TypeError, 

231 ModuleVisibilityError, 

232 UnresolvableDependencyError, 

233 ): 

234 logger.debug("rag_working_memory_not_available") 

235 

236 if working_memory is not None: 

237 self._working_memory = working_memory 

238 

239 # Optional: graph store 

240 graph_store_available = False 

241 try: 

242 from lexigram.contracts.data.graph.protocols import GraphStoreProtocol 

243 

244 graph_store = await container.resolve_optional(GraphStoreProtocol) 

245 graph_store_available = graph_store is not None 

246 except ( 

247 LookupError, 

248 RuntimeError, 

249 AttributeError, 

250 ImportError, 

251 TypeError, 

252 ModuleVisibilityError, 

253 UnresolvableDependencyError, 

254 ): 

255 pass 

256 

257 logger.info( 

258 "rag_provider_booted", 

259 working_memory=working_memory is not None, 

260 graph_store=graph_store_available, 

261 ) 

262 

263 async def shutdown(self) -> None: 

264 pass 

265 

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

267 """Check RAG provider health — verifies embedding service and vector store. 

268 

269 Returns: 

270 :class:`~lexigram.contracts.core.HealthCheckResult` with status 

271 ``healthy`` when all configured dependencies are reachable, or 

272 ``degraded``/``unhealthy`` otherwise. 

273 """ 

274 details: dict = { 

275 "embedding_service": "unconfigured", 

276 "vector_store": "unconfigured", 

277 } 

278 overall = HealthStatus.HEALTHY 

279 

280 container: BootContainerProtocol | None = getattr( 

281 self, "_booted_container", None 

282 ) 

283 if container is None: 

284 return HealthCheckResult( 

285 component="rag", 

286 status=HealthStatus.HEALTHY, 

287 details={"note": "RAG provider not yet booted"}, 

288 ) 

289 

290 # Check embedding client if available 

291 try: 

292 from lexigram.contracts.ai import EmbeddingClientProtocol 

293 

294 embedding_client = await container.resolve_optional(EmbeddingClientProtocol) 

295 if embedding_client is not None and hasattr( 

296 embedding_client, "health_check" 

297 ): 

298 emb_result = await embedding_client.health_check(timeout=timeout) 

299 if hasattr(emb_result, "status"): 

300 details["embedding_service"] = emb_result.status 

301 if emb_result.status != HealthStatus.HEALTHY: 

302 overall = HealthStatus.DEGRADED 

303 else: 

304 details["embedding_service"] = "ok" 

305 else: 

306 details["embedding_service"] = "not_configured" 

307 except (LookupError, RuntimeError, AttributeError) as exc: 

308 details["embedding_service"] = f"error: {exc}" 

309 overall = HealthStatus.DEGRADED 

310 

311 # Check vector store if available 

312 try: 

313 from lexigram.contracts.ai import DocumentVectorStoreProtocol 

314 

315 vector_store = await container.resolve_optional(DocumentVectorStoreProtocol) 

316 if vector_store is not None and hasattr(vector_store, "health_check"): 

317 vs_result = await vector_store.health_check(timeout=timeout) 

318 if hasattr(vs_result, "status"): 

319 details["vector_store"] = vs_result.status 

320 if vs_result.status != HealthStatus.HEALTHY: 

321 overall = HealthStatus.DEGRADED 

322 else: 

323 details["vector_store"] = "ok" 

324 else: 

325 details["vector_store"] = "not_configured" 

326 except (LookupError, RuntimeError, AttributeError) as exc: 

327 details["vector_store"] = f"error: {exc}" 

328 overall = HealthStatus.DEGRADED 

329 

330 return HealthCheckResult( 

331 component="rag", 

332 status=overall, 

333 details=details, 

334 ) 

335 

336 

337__all__ = ["RAGProvider"]