Coverage for llm_dataset_engine/stages/pipeline_stage.py: 41%

1"""

2 Base pipeline stage abstraction.

3

4 Defines the contract for all processing stages using Template Method

5 pattern for execution flow.

6"""

7

8 from abc import ABC, abstractmethod

9 from typing import Any, Generic, TypeVar

10

11 from llm_dataset_engine.core.models import CostEstimate, ValidationResult

12 from llm_dataset_engine.utils import get_logger

13

14# Type variables for input and output

15 TInput = TypeVar("TInput")

16 TOutput = TypeVar("TOutput")

17

18 logger = get_logger(__name__)

19

20

21 class PipelineStage(ABC, Generic[TInput, TOutput]):

22 """

23 Abstract base class for all pipeline stages.

24

25 Implements Template Method pattern with hooks for extensibility.

26 All stages follow Single Responsibility and are composable.

27 """

28

29 def __init__(self, name: str):

30 """

31 Initialize pipeline stage.

32

33 Args:

34 name: Human-readable stage name

35 """

36 self.name = name

37 self.logger = get_logger(f"{__name__}.{name}")

38

39 @abstractmethod

40 def process(self, input_data: TInput, context: Any) -> TOutput:

41 """

42 Core processing logic (must be implemented by subclasses).

43

44 Args:

45 input_data: Input data for this stage

46 context: Execution context with shared state

47

48 Returns:

49 Processed output data

50 """

51 pass

52

53 @abstractmethod

54 def validate_input(self, input_data: TInput) -> ValidationResult:

55 """

56 Validate input before processing.

57

58 Args:

59 input_data: Input to validate

60

61 Returns:

62 ValidationResult with errors/warnings

63 """

64 pass

65

66 def execute(self, input_data: TInput, context: Any) -> TOutput:

67 """

68 Execute stage with pre/post hooks (Template Method).

69

70 This method orchestrates the execution flow and should not

71 be overridden.

72

73 Args:

74 input_data: Input data

75 context: Execution context

76

77 Returns:

78 Processed output

79

80 Raises:

81 ValueError: If input validation fails

82 """

83 self.logger.info(f"Starting stage: {self.name}")

84

85 # Pre-processing hook

86 self.before_process(context)

87

88 # Validate input

89 validation = self.validate_input(input_data)

90 if not validation.is_valid:

91 error_msg = f"Input validation failed: {validation.errors}"

92 self.logger.error(error_msg)

93 raise ValueError(error_msg)

94

95 if validation.warnings:

96 for warning in validation.warnings:

97 self.logger.warning(warning)

98

99 # Core processing

100 try:

101 result = self.process(input_data, context)

102 self.logger.info(f"Completed stage: {self.name}")

103

104 # Post-processing hook

105 self.after_process(result, context)

106

107 return result

108 except Exception as e:

109 self.logger.error(f"Stage {self.name} failed: {e}")

110 error_decision = self.on_error(e, context)

111 raise error_decision

112

113 def before_process(self, context: Any) -> None:

114 """

115 Hook called before processing (default: no-op).

116

117 Args:

118 context: Execution context

119 """

120 pass

121

122 def after_process(self, result: TOutput, context: Any) -> None:

123 """

124 Hook called after successful processing (default: no-op).

125

126 Args:

127 result: Processing result

128 context: Execution context

129 """

130 pass

131

132 def on_error(self, error: Exception, context: Any) -> Exception:

133 """

134 Hook called on processing error (default: re-raise).

135

136 Args:

137 error: The exception that occurred

138 context: Execution context

139

140 Returns:

141 Exception to raise (can transform error)

142 """

143 return error

144

145 @abstractmethod

146 def estimate_cost(self, input_data: TInput) -> CostEstimate:

147 """

148 Estimate processing cost for this stage.

149

150 Args:

151 input_data: Input data to estimate for

152

153 Returns:

154 Cost estimate

155 """

156 pass

157