Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/pipeline/builder.py: 24%

1from __future__ import annotations 

2 

3from pathlib import Path 

4from typing import TYPE_CHECKING, Any 

5 

6import yaml 

7 

8from lexigram.ai.rag.config import ( 

9 PipelineConfig, 

10 PipelineStageType, 

11) 

12from lexigram.ai.rag.exceptions import MissingCitationsError 

13from lexigram.ai.rag.pipeline._stage_factory import build_pipeline_stages 

14from lexigram.ai.rag.pipeline.base import PipelineStageProtocol 

15from lexigram.ai.rag.pipeline.executor import PipelineExecutor 

16from lexigram.ai.rag.pipeline.types import ErrorStrategy 

17from lexigram.contracts.ai.exceptions import RAGError 

18from lexigram.contracts.ai.rag import RAGContext, RAGEvaluatorProtocol, RAGResponse 

19from lexigram.logging import ( 

20 get_logger, 

21) 

22from lexigram.result import Err, Ok, Result 

23 

24if TYPE_CHECKING: 

25 from lexigram.ai.rag.pipeline.types import PipelineContext 

26 from lexigram.contracts.ai.memory import WorkingMemoryProtocol 

27 

28from lexigram.primitives.builder import AbstractBuilder 

29 

30logger = get_logger(__name__) 

31 

32 

33class RAGPipeline: 

34 """Main RAG pipeline that orchestrates all stages. 

35 

36 This class provides a simple interface for executing the complete 

37 RAG pipeline with configurable stages and error handling. 

38 """ 

39 

40 def __init__( 

41 self, 

42 config: PipelineConfig, 

43 stages: list[PipelineStageProtocol], 

44 evaluator: RAGEvaluatorProtocol | None = None, 

45 working_memory: WorkingMemoryProtocol | None = None, 

46 ): 

47 """Initialize the RAG pipeline. 

48 

49 Args: 

50 config: Pipeline configuration 

51 stages: List of pipeline stages 

52 evaluator: Optional evaluator implementing 

53 :class:`~lexigram.contracts.ai.rag.RAGEvaluatorProtocol` for 

54 automatic per-request quality evaluation. Evaluation 

55 frequency is controlled by :attr:`~lexigram.ai.rag.config.PipelineConfig.auto_evaluate_every_n`. 

56 working_memory: Optional working memory for context enrichment. 

57 """ 

58 self.config = config 

59 self.stages = stages 

60 self.evaluator = evaluator 

61 self._working_memory = working_memory 

62 self._request_count: int = 0 

63 error_strategy = config.default_error_strategy 

64 if not isinstance(error_strategy, ErrorStrategy): 

65 if isinstance(error_strategy, str): 

66 error_strategy = ErrorStrategy(error_strategy) 

67 else: 

68 error_strategy = ErrorStrategy.GRACEFUL 

69 

70 self.executor = PipelineExecutor( 

71 stages=stages, 

72 default_error_strategy=error_strategy, 

73 max_retries=config.max_retries, 

74 retry_delay=config.retry_delay, 

75 ) 

76 

77 async def run( 

78 self, 

79 query: str, 

80 documents: list[str] | None = None, 

81 document_paths: list[str] | None = None, 

82 metadata: dict[str, Any] | None = None, 

83 ) -> PipelineContext: 

84 """Execute the RAG pipeline. 

85 

86 Args: 

87 query: User query 

88 documents: Optional list of document content strings 

89 document_paths: Optional list of document file paths 

90 metadata: Optional custom metadata 

91 

92 Returns: 

93 Pipeline context with results 

94 """ 

95 from lexigram.ai.rag.pipeline.types import PipelineContext 

96 

97 # Create initial context 

98 context = PipelineContext( 

99 query=query, 

100 documents=documents or [], 

101 document_paths=document_paths or [], 

102 metadata=metadata or {}, 

103 ) 

104 

105 # Enrich query with memory context if working memory is available 

106 if self._working_memory: 

107 try: 

108 memory_entries = await self._working_memory.assemble( 

109 query=query, 

110 token_budget=1024, 

111 owner_id=(metadata or {}).get("user_id") 

112 or (metadata or {}).get("session_id") 

113 or "anonymous", 

114 ) 

115 if memory_entries: 

116 memory_context = "\n".join( 

117 f"[{e.role}]: {e.content}" for e in memory_entries 

118 ) 

119 context.metadata["memory_context"] = memory_context 

120 context.metadata["memory_entries_count"] = len(memory_entries) 

121 except (RuntimeError, TypeError, AttributeError, LookupError) as e: 

122 logger.warning("rag_memory_enrichment_failed", error=str(e)) 

123 

124 # Execute pipeline 

125 context = await self.executor.execute(context) 

126 

127 # Citation enforcement (P7.2) — fail fast before incrementing counters 

128 if ( 

129 self.config.require_citations 

130 and context.synthesis_result is not None 

131 and not context.synthesis_result.citations 

132 ): 

133 msg = ( 

134 f"Pipeline '{self.config.name}' requires citations but none were " 

135 "produced by the synthesis stage." 

136 ) 

137 raise MissingCitationsError(msg) 

138 

139 # Auto-evaluation hook (P7.1) 

140 self._request_count += 1 

141 if ( 

142 self.evaluator is not None 

143 and self.config.auto_evaluate_every_n is not None 

144 and self._request_count % self.config.auto_evaluate_every_n == 0 

145 and context.synthesis_result is not None 

146 ): 

147 docs = context.optimized_chunks or context.retrieved_chunks 

148 try: 

149 report = await self.evaluator.evaluate( 

150 query=context.query, 

151 retrieved_docs=docs, 

152 generated_answer=context.synthesis_result.response, 

153 ) 

154 context.metadata["evaluation_report"] = report 

155 except (RuntimeError, TypeError, AttributeError, OSError, ValueError): 

156 logger.warning( 

157 "Auto-evaluation failed for request %s", 

158 context.request_id, 

159 exc_info=True, 

160 ) 

161 

162 return context 

163 

164 async def execute(self, context: RAGContext) -> Result[RAGResponse, RAGError]: 

165 """Execute the RAG pipeline per the contract protocol. 

166 

167 Args: 

168 context: Pipeline context with query and optional config/filters. 

169 

170 Returns: 

171 Ok(RAGResponse) on success, Err(RAGError) on failure. 

172 """ 

173 try: 

174 pipeline_ctx = await self.run( 

175 query=context.query, 

176 metadata=context.filters, 

177 ) 

178 synthesis = pipeline_ctx.synthesis_result 

179 answer = synthesis.response if synthesis is not None else "" 

180 citations = synthesis.citations if synthesis is not None else None 

181 response = RAGResponse( 

182 answer=answer, 

183 sources=[], 

184 citations=citations, 

185 ) 

186 return Ok(response) 

187 except RAGError as exc: 

188 return Err(exc) 

189 

190 async def run_parallel( 

191 self, 

192 query: str, 

193 stages: list[PipelineStageProtocol] | None = None, 

194 **kwargs: Any, 

195 ) -> PipelineContext: 

196 """Execute pipeline stages in parallel. 

197 

198 Args: 

199 query: User query 

200 stages: Stages to execute in parallel (default: all stages) 

201 **kwargs: Additional context parameters 

202 

203 Returns: 

204 Pipeline context with results 

205 """ 

206 from lexigram.ai.rag.pipeline.types import PipelineContext 

207 

208 context = PipelineContext(query=query, **kwargs) 

209 return await self.executor.execute_parallel(context, stages) 

210 

211 

212class PipelineBuilder(AbstractBuilder[RAGPipeline]): 

213 """Builder for constructing RAG pipelines with fluent API. 

214 

215 The builder provides a convenient way to configure and build 

216 pipelines programmatically or from configuration files. 

217 """ 

218 

219 def __init__(self) -> None: 

220 """Initialize the pipeline builder.""" 

221 super().__init__() 

222 self.config = PipelineConfig() 

223 self._custom_stages: list[PipelineStageProtocol] = [] 

224 

225 def with_name(self, name: str) -> PipelineBuilder: 

226 """Set pipeline name. 

227 

228 Args: 

229 name: Pipeline name 

230 

231 Returns: 

232 Self for chaining 

233 """ 

234 self.config.name = name 

235 return self 

236 

237 def with_description(self, description: str) -> PipelineBuilder: 

238 """Set pipeline description. 

239 

240 Args: 

241 description: Pipeline description 

242 

243 Returns: 

244 Self for chaining 

245 """ 

246 self.config.description = description 

247 return self 

248 

249 def with_ingestion(self, **kwargs: Any) -> PipelineBuilder: 

250 """Configure ingestion stage. 

251 

252 Args: 

253 **kwargs: Ingestion configuration parameters 

254 

255 Returns: 

256 Self for chaining 

257 """ 

258 for key, value in kwargs.items(): 

259 if hasattr(self.config.ingestion, key): 

260 setattr(self.config.ingestion, key, value) 

261 return self 

262 

263 def with_query_processing(self, **kwargs: Any) -> PipelineBuilder: 

264 """Configure query processing stage. 

265 

266 Args: 

267 **kwargs: Query processing configuration parameters 

268 

269 Returns: 

270 Self for chaining 

271 """ 

272 for key, value in kwargs.items(): 

273 if hasattr(self.config.query_processing, key): 

274 setattr(self.config.query_processing, key, value) 

275 return self 

276 

277 def with_retrieval(self, **kwargs: Any) -> PipelineBuilder: 

278 """Configure retrieval stage. 

279 

280 Args: 

281 **kwargs: Retrieval configuration parameters 

282 

283 Returns: 

284 Self for chaining 

285 """ 

286 for key, value in kwargs.items(): 

287 if hasattr(self.config.retrieval, key): 

288 setattr(self.config.retrieval, key, value) 

289 return self 

290 

291 def with_context_optimization(self, **kwargs: Any) -> PipelineBuilder: 

292 """Configure context optimization stage. 

293 

294 Args: 

295 **kwargs: Context optimization configuration parameters 

296 

297 Returns: 

298 Self for chaining 

299 """ 

300 for key, value in kwargs.items(): 

301 if hasattr(self.config.context_optimization, key): 

302 setattr(self.config.context_optimization, key, value) 

303 return self 

304 

305 def with_synthesis(self, **kwargs: Any) -> PipelineBuilder: 

306 """Configure synthesis stage. 

307 

308 Args: 

309 **kwargs: Synthesis configuration parameters 

310 

311 Returns: 

312 Self for chaining 

313 """ 

314 for key, value in kwargs.items(): 

315 if hasattr(self.config.synthesis, key): 

316 setattr(self.config.synthesis, key, value) 

317 return self 

318 

319 def with_quality_assurance(self, **kwargs: Any) -> PipelineBuilder: 

320 """Configure quality assurance stage. 

321 

322 Args: 

323 **kwargs: Quality assurance configuration parameters 

324 

325 Returns: 

326 Self for chaining 

327 """ 

328 for key, value in kwargs.items(): 

329 if hasattr(self.config.quality_assurance, key): 

330 setattr(self.config.quality_assurance, key, value) 

331 return self 

332 

333 def with_post_processing(self, **kwargs: Any) -> PipelineBuilder: 

334 """Configure post-processing stage. 

335 

336 Args: 

337 **kwargs: Post-processing configuration parameters 

338 

339 Returns: 

340 Self for chaining 

341 """ 

342 for key, value in kwargs.items(): 

343 if hasattr(self.config.post_processing, key): 

344 setattr(self.config.post_processing, key, value) 

345 return self 

346 

347 def with_error_strategy( 

348 self, 

349 strategy: ErrorStrategy, 

350 max_retries: int = 3, 

351 retry_delay: float = 1.0, 

352 ) -> PipelineBuilder: 

353 """Configure global error handling. 

354 

355 Args: 

356 strategy: Default error handling strategy 

357 max_retries: Maximum number of retries 

358 retry_delay: Initial delay between retries 

359 

360 Returns: 

361 Self for chaining 

362 """ 

363 self.config.default_error_strategy = strategy 

364 self.config.max_retries = max_retries 

365 self.config.retry_delay = retry_delay 

366 return self 

367 

368 def with_stages(self, stages: list[PipelineStageType]) -> PipelineBuilder: 

369 """Set the ordered list of pipeline stages. 

370 

371 Args: 

372 stages: List of stages 

373 

374 Returns: 

375 Self for chaining 

376 """ 

377 self.config.stages = stages 

378 return self 

379 

380 def with_custom_stage(self, stage: PipelineStageProtocol) -> PipelineBuilder: 

381 """Add a custom pipeline stage. 

382 

383 Args: 

384 stage: Custom pipeline stage 

385 

386 Returns: 

387 Self for chaining 

388 """ 

389 self._custom_stages.append(stage) 

390 return self 

391 

392 def from_dict(self, config_dict: dict[str, Any]) -> PipelineBuilder: 

393 """Load configuration from dictionary. 

394 

395 Args: 

396 config_dict: Configuration dictionary 

397 

398 Returns: 

399 Self for chaining 

400 """ 

401 self.config = PipelineConfig.from_dict(config_dict) 

402 return self 

403 

404 async def from_yaml(self, yaml_path: str | Path) -> PipelineBuilder: 

405 """Load configuration from YAML file. 

406 

407 Args: 

408 yaml_path: Path to YAML configuration file 

409 

410 Returns: 

411 Self for chaining 

412 """ 

413 import asyncio 

414 

415 def _load_yaml() -> Any: 

416 with open(yaml_path) as f: 

417 return yaml.safe_load(f.read()) 

418 

419 config_dict = await asyncio.to_thread(_load_yaml) 

420 return self.from_dict(config_dict) 

421 

422 # ----------------------------------------------------------------- 

423 # High-level convenience API (Phase 10 DX) 

424 # ----------------------------------------------------------------- 

425 

426 def retrieve( 

427 self, 

428 strategy: str = "hybrid", 

429 top_k: int = 10, 

430 **kwargs: Any, 

431 ) -> PipelineBuilder: 

432 """Configure retrieval with high-level parameters. 

433 

434 Args: 

435 strategy: Retrieval strategy name (``"hybrid"``, ``"dense"``, 

436 ``"sparse"``). 

437 top_k: Number of chunks to retrieve. 

438 **kwargs: Additional retrieval parameters. 

439 """ 

440 self.config.retrieval.enabled = True 

441 self.config.retrieval.strategy = strategy 

442 self.config.retrieval.top_k = top_k 

443 for key, value in kwargs.items(): 

444 if hasattr(self.config.retrieval, key): 

445 setattr(self.config.retrieval, key, value) 

446 return self 

447 

448 def rerank( 

449 self, 

450 strategy: str = "cross-encoder", 

451 top_k: int = 5, 

452 **kwargs: Any, 

453 ) -> PipelineBuilder: 

454 """Configure context optimization / reranking. 

455 

456 Args: 

457 strategy: Reranking strategy name. 

458 top_k: Number of chunks to keep after reranking. 

459 **kwargs: Additional reranking parameters. 

460 """ 

461 self.config.context_optimization.enabled = True 

462 self.config.context_optimization.strategy = strategy 

463 self.config.context_optimization.top_k = top_k 

464 for key, value in kwargs.items(): 

465 if hasattr(self.config.context_optimization, key): 

466 setattr(self.config.context_optimization, key, value) 

467 return self 

468 

469 def synthesize( 

470 self, 

471 strategy: str = "abstractive", 

472 model: str | None = None, 

473 **kwargs: Any, 

474 ) -> PipelineBuilder: 

475 """Configure synthesis with high-level parameters. 

476 

477 Args: 

478 strategy: Synthesis strategy (``"abstractive"``, ``"extractive"``). 

479 model: LLM model identifier for generation. 

480 **kwargs: Additional synthesis parameters. 

481 """ 

482 self.config.synthesis.enabled = True 

483 self.config.synthesis.strategy = strategy # type: ignore[assignment] 

484 if model is not None: 

485 self.config.synthesis.model = model 

486 for key, value in kwargs.items(): 

487 if hasattr(self.config.synthesis, key): 

488 setattr(self.config.synthesis, key, value) 

489 return self 

490 

491 def with_citations(self, required: bool = True) -> PipelineBuilder: 

492 """Enable citation tracking in the pipeline. 

493 

494 Args: 

495 required: Whether citations are required (pipeline fails without 

496 them if set to ``True``). 

497 """ 

498 self.config.require_citations = required 

499 return self 

500 

501 def with_evaluation( 

502 self, 

503 evaluator: RAGEvaluatorProtocol | None = None, 

504 every_n: int = 1, 

505 ) -> PipelineBuilder: 

506 """Enable automatic evaluation of pipeline outputs. 

507 

508 Args: 

509 evaluator: Optional pre-built evaluator instance. 

510 every_n: Evaluate every *n*-th request (default: every request). 

511 """ 

512 self._evaluator = evaluator 

513 self.config.auto_evaluate_every_n = every_n 

514 return self 

515 

516 def with_timeout(self, **stage_timeouts: float) -> PipelineBuilder: 

517 """Set per-stage timeouts (in seconds). 

518 

519 Keyword arguments map stage names to their timeout values:: 

520 

521 builder.with_timeout(retrieval=10.0, synthesis=30.0) 

522 

523 Args: 

524 **stage_timeouts: Mapping of stage name to timeout in seconds. 

525 """ 

526 if not hasattr(self.config, "stage_timeouts"): 

527 setattr(self.config, "stage_timeouts", {}) # noqa: B010 

528 self.config.stage_timeouts.update(stage_timeouts) # type: ignore[attr-defined] 

529 return self 

530 

531 def with_working_memory( 

532 self, 

533 memory: WorkingMemoryProtocol, 

534 ) -> PipelineBuilder: 

535 """Attach working memory for context enrichment. 

536 

537 Args: 

538 memory: Working memory instance. 

539 """ 

540 self._working_memory = memory 

541 return self 

542 

543 def build(self) -> RAGPipeline: 

544 """Build the RAG pipeline. 

545 

546 Returns: 

547 Configured RAG pipeline 

548 

549 Raises: 

550 ValueError: If configuration is invalid 

551 """ 

552 # Validate configuration 

553 self._validate_config() 

554 

555 # Build stages 

556 stages = self._build_stages() 

557 

558 # Add custom stages 

559 stages.extend(self._custom_stages) 

560 

561 # Create pipeline 

562 return RAGPipeline( 

563 config=self.config, 

564 stages=stages, 

565 evaluator=getattr(self, "_evaluator", None), 

566 working_memory=getattr(self, "_working_memory", None), 

567 ) 

568 

569 def _validate_config(self) -> None: 

570 """Validate pipeline configuration. 

571 

572 Raises: 

573 ValueError: If configuration is invalid 

574 """ 

575 # Check that at least synthesis is enabled 

576 if not self.config.synthesis.enabled: 

577 msg = "Synthesis must be enabled in pipeline configuration" 

578 raise ValueError(msg) 

579 

580 # Check that retrieval is enabled if synthesis is enabled 

581 if self.config.synthesis.enabled and not self.config.retrieval.enabled: 

582 raise ValueError( 

583 "Retrieval stage must be enabled when synthesis is enabled", 

584 ) 

585 

586 # Validate quality thresholds 

587 qa_config = self.config.quality_assurance 

588 if qa_config.enabled: 

589 if not (0 <= qa_config.min_faithfulness <= 1): 

590 msg = "qa_config.min_faithfulness must be between 0 and 1" 

591 raise ValueError(msg) 

592 if not (0 <= qa_config.min_relevance <= 1): 

593 msg = "qa_config.min_relevance must be between 0 and 1" 

594 raise ValueError(msg) 

595 if not (0 <= qa_config.min_confidence <= 1): 

596 msg = "qa_config.min_confidence must be between 0 and 1" 

597 raise ValueError(msg) 

598 

599 def _build_stages(self) -> list[PipelineStageProtocol]: 

600 """Build pipeline stages based on configuration. 

601 

602 Returns: 

603 List of configured pipeline stages. 

604 """ 

605 return build_pipeline_stages(self.config)