1"""Factory utilities for constructing configured RAG pipeline stages."""
2
3from __future__ import annotations
4
5from lexigram.ai.rag.config import PipelineConfig, PipelineStageType
6from lexigram.ai.rag.pipeline.base import PipelineStageProtocol
7
8
9def build_pipeline_stages(config: PipelineConfig) -> list[PipelineStageProtocol]:
10 """Build pipeline stages based on the provided pipeline config."""
11 from lexigram.ai.rag.pipeline.stages import (
12 QualityAssuranceStage,
13 RetrievalStage,
14 SynthesisStage,
15 )
16
17 stages: list[PipelineStageProtocol] = []
18
19 for stage_type in config.stages:
20 if stage_type == PipelineStageType.RETRIEVAL and config.retrieval.enabled:
21 stages.append(RetrievalStage(config=config.retrieval))
22 elif stage_type == PipelineStageType.SYNTHESIS and config.synthesis.enabled:
23 stages.append(SynthesisStage(config=config.synthesis))
24 elif (
25 stage_type == PipelineStageType.QUALITY_ASSURANCE
26 and config.quality_assurance.enabled
27 ):
28 stages.append(QualityAssuranceStage(config=config.quality_assurance))
29
30 return stages